import librosa
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.multiclass import OneVsRestClassifier
from sklearn.preprocessing import StandardScaler
import sys


# 加载数据集
def load_data():
    # 加载 tbb、aichen、xsc 三种乐器的音频数据
    # 使用原始字符串以避免转义字符问题
    tbb, sr1 = librosa.load(sys.argv[1])
    aichen, sr2 = librosa.load(sys.argv[2])
    xsc, sr3 = librosa.load(sys.argv[3])

    # 提取 MFCC 特征，这里也就是不同人声音音色提取
    tbb_mfcc = librosa.feature.mfcc(y=tbb, sr=sr1)
    aichen_mfcc = librosa.feature.mfcc(y=aichen, sr=sr2)
    xsc_mfcc = librosa.feature.mfcc(y=xsc, sr=sr3)

    # 将不同人声音色的 MFCC 特征合并成一个数据集
    X = np.concatenate((tbb_mfcc.T, aichen_mfcc.T, xsc_mfcc.T), axis=0)

    # 生成标签向量
    # 正确使用 np.concatenate 函数，将数组作为元组传递
    y = np.concatenate((np.zeros(len(tbb_mfcc.T)), np.ones(len(aichen_mfcc.T)), 2 * np.ones(len(xsc_mfcc.T))), axis=0)


    return X, y


# 训练模型
def train(X, y):
    # 将数据集分成训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # 对特征数据进行标准化处理
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    # 使用逻辑回归算法进行多类别分类
    # 使用 OneVsRestClassifier 包装 LogisticRegression
    # 增加 max_iter 参数，以增加迭代次数
    model = OneVsRestClassifier(LogisticRegression(max_iter=1000))


    # 训练模型
    model.fit(X_train_scaled, y_train)

    return model, scaler


# 测试模型
def predict(model, scaler, audio_file):
    # 加载音频文件并提取 MFCC 特征
    y, sr = librosa.load(audio_file)
    mfcc = librosa.feature.mfcc(y=y, sr=sr)
    mfcc_scaled = scaler.transform(mfcc.T)

    # 进行多类别分类预测
    label = model.predict(mfcc_scaled)
    proba = model.predict_proba(mfcc_scaled)

    # 获取概率最大的类别标签
    # 修正获取最大概率类别索引的方式
    max_prob_idx = np.argmax(proba)
    max_prob_label = np.argmax(proba[max_prob_idx])


    return max_prob_label


if __name__ == '__main__':
    try:
        # 加载数据集
        X, y = load_data()
        # 训练模型
        model, scaler = train(X, y)
        # 测试模型
        # 使用原始字符串以避免转义字符问题
        label = predict(model, scaler, sys.argv[4])
        print(label)
    except Exception as e:
        import traceback
        traceback.print_exc()
